# -*- coding: utf-8 -*-
"""
A股均值回归策略
包含沪深300成分股筛选、涨跌停处理、T+1交易限制等A股特有机制
"""

from zipline.api import order_target_percent, order_target, symbol, get_datetime
from zipline.finance import commission, slippage
from strategies.utils.a_share_utils import get_csi300_constituents, is_price_limit, apply_t1_restriction

def initialize(context):
    # 初始化持仓的最高市值记录
    context.portfolio.positions.high_watermark = {}
    """
    初始化策略
    """
    # 设置策略参数
    context.lookback = 20  # 均线周期
    context.std_threshold = 2.0  # 偏离阈值
    context.max_position = 0.05  # 单票最大仓位
    context.stop_loss = 0.96  # 止损阈值 (4%)
    context.take_profit = 1.08  # 止盈阈值 (8%)
    context.trailing_stop = 0.98  # 移动止盈回撤 (2%)
    
    # 设置股票池为空，将在首次运行时初始化
    context.index_components = []
    
    # 设置再平衡频率 (每天运行)
    context.rebalance_frequency = 1
    context.days_until_rebalance = 0
    
    # 设置交易费用 (A股标准)
    context.set_commission(commission.PerShare(cost=0.0003, min_trade_cost=5))
    context.set_slippage(slippage.FixedSlippage(spread=0.001))

def handle_data(context, data):
    """
    每日交易逻辑
    """
    # 检查再平衡日
    context.days_until_rebalance -= 1
    if context.days_until_rebalance <= 0:
        # 更新股票池 (季度再平衡)
        context.index_components = get_csi300_constituents(context, data)
        context.days_until_rebalance = 63  # 约3个月
    
    # 首次运行时初始化股票池
    if not context.index_components:
        context.index_components = get_csi300_constituents(context, data)
    
    # 执行均值回归策略
    for stock in context.index_components:
        # 初始化持仓的最高市值
        if stock not in context.portfolio.positions.high_watermark:
            context.portfolio.positions.high_watermark[stock] = 0
        # 跳过涨跌停股票
        if is_price_limit(stock, data):
            continue
            
        # 获取价格数据
        prices = data.history(stock, 'close', context.lookback, '1d')
        current_price = data.current(stock, 'close')
        
        # 计算均线和标准差
        ma = prices.mean()
        std = prices.std()
        
        # 生成交易信号
        position = context.portfolio.positions[stock]
        
        # 买入信号：价格低于均值-2倍标准差
        if current_price < ma - context.std_threshold * std:
            # 应用T+1限制
            closeable_amount = apply_t1_restriction(context, stock)
            if closeable_amount == 0:  # 无持仓时可买入
                order_target_percent(stock, context.max_position)
        
        # 持仓管理
        if position.amount > 0:
            cost_basis = position.cost_basis
            current_value = position.last_sale_price
            
            # 止损逻辑
            if current_value < cost_basis * context.stop_loss:
                order_target(stock, 0)
            
            # 移动止盈逻辑
            elif current_value > cost_basis * context.take_profit:
                if current_value < context.portfolio.positions.high_watermark[stock] * context.trailing_stop:
                    order_target(stock, 0)
                else:
                    context.portfolio.positions.high_watermark[stock] = max(
                        context.portfolio.positions.high_watermark[stock], current_value
                    )

def analyze(context, perf):
    """
    回测结果分析
    """
    # 输出策略表现
    print("策略年化收益率: %.2f%%" % (perf.returns.mean() * 252 * 100))
    print("最大回撤: %.2f%%" % (perf.max_drawdown() * 100))
    print("夏普比率: %.2f" % perf.sharpe_ratio)
    
    # 绘制净值曲线
    perf.portfolio_value.plot()